import httpx
from fastapi import FastAPI, Request
from starlette.responses import StreamingResponse


class AppLogger:
    '''
    写日志
    '''
    def __init__(self, log_file="llm.log"):
        """Initialize the logger with a file that will be cleared on startup."""
        self.log_file = log_file
        # Clear the log file on startup
        with open(self.log_file, 'w') as f:
            f.write("")

    def log(self, message):
        """Log a message to both file and console.
        message: 要写的日志内容
        """

        # Log to file
        with open(self.log_file, 'a') as f:
            f.write(message + "\n")

        # Log to console
        print(message)


app = FastAPI(title="LLM API Logger")  # web开发框架，用于后续定义POST接口
logger = AppLogger("llm.log")


# 用app定义了一个POST接口
@app.post("/chat/completions")
async def proxy_request(request: Request):
    '''
    request: HTTP请求对象，包含了请求的所有信息
    '''

    # 1.获取cline请求数据
    body_bytes = await request.body()
    body_str = body_bytes.decode('utf-8')
    logger.log(f"模型请求：{body_str}")
    body = await request.json()

    # 2.处理模型返回的代码
    logger.log("模型返回：\n")
    # 定义一个异步生成器函数，用于处理事件流
    async def event_stream():
        # 创建一个异步HTTP客户端，设置无超时限制（None表示不超时）
        # 使用async with确保客户端使用完毕后正确关闭资源
        async with httpx.AsyncClient(timeout=None) as client:
            # 3.向openrouter.ai发送请求
            async with client.stream(
            # 使用客户端的stream方法发起POST请求并处理流式响应
                    "POST",
                    "https://openrouter.ai/api/v1/chat/completions",
                    json=body,
                    headers={
                        "Content-Type": "application/json",
                        "Accept": "text/event-stream",
                        "Authorization": request.headers.get("Authorization"),
                    },
            ) as response:
                # 异步迭代响应的每一行内容
                async for line in response.aiter_lines():
                    # 将获取到的行记录到日志
                    logger.log(line)
                    # 通过生成器将行内容返回，末尾添加换行符确保格式正确
                    yield f"{line}\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
